#!/bin/bash

# Copyright     2017  David Snyder
#               2017  Johns Hopkins University (Author: Daniel Povey)
#               2017  Johns Hopkins University (Author: Daniel Garcia Romero)
# Apache 2.0.

# This script extracts embeddings (called "dvectors" here) from a set of
# utterances, given features and a trained DNN.  The purpose of this script
# is analogous to sid/extract_ivectors.sh: it creates archives of
# vectors that are used in speaker recognition.  Like ivectors, dvectors can
# be used in PLDA or a similar backend for scoring.

# Begin configuration section.
nj=30
cmd="run.pl"

apply_cmvn_sliding=true
apply_cmvn_utt=false

stage=0

echo "$0 $@"  # Print the command line for logging

if [ -f path.sh ]; then . ./path.sh; fi
. parse_options.sh || exit 1;

if [ $# != 3 ]; then
  echo "Usage: $0 <nnet-dir> <data> <dvector-dir>"
  echo " e.g.: $0 exp/dvector_nnet data/train exp/dvectors_train"
  echo "main options (for others, see top of script file)"
  echo "  --config <config-file>                           # config containing options"
  echo "  --cmd (utils/run.pl|utils/queue.pl <queue opts>) # how to run jobs."
  echo "  --use-gpu <bool|false>                           # If true, use GPU."
  echo "  --nj <n|30>                                      # Number of jobs"
  echo "  --stage <stage|0>                                # To control partial reruns"
fi

srcdir=$1
data=$2
dir=$3

for f in $srcdir/final.raw $data/feats.scp $data/vad.scp ; do
  [ ! -f $f ] && echo "No such file $f" && exit 1;
done

mkdir -p $dir/log
cp $data/{spk2utt,utt2spk} $dir/

nnet=$srcdir/final.raw
if [ -f $srcdir/extract.config ] ; then
  echo "$0: using $srcdir/extract.config to extract dvectors"
  nnet="nnet3-copy --nnet-config=$srcdir/extract.config $srcdir/final.raw - |"
fi

utils/split_data.sh $data $nj
echo "$0: extracting dvectors for $data"
sdata=$data/split$nj/JOB

# Set up the features
if $apply_cmvn_sliding; then
  feat="ark:apply-cmvn-sliding --norm-vars=false --center=true --cmn-window=300 scp:${sdata}/feats.scp ark:- | select-voiced-frames ark:- scp,s,cs:${sdata}/vad.scp ark:- |"
elif $apply_cmvn_utt; then
  feat="ark:apply-cmvn --norm-means=true --norm-vars=false scp:${sdata}/cmvn.scp scp:${sdata}/feats.scp ark:- | select-voiced-frames ark:- scp,s,cs:${sdata}/vad.scp ark:- |"
else
  feat="ark:copy-feats scp:${sdata}/feats.scp ark:- | select-voiced-frames ark:- scp,s,cs:${sdata}/vad.scp ark:- |"
fi

if [ $stage -le 0 ]; then
  echo "$0: extracting dvectors from nnet"
  $cmd JOB=1:$nj ${dir}/log/extract.JOB.log \
    nnet3-compute --use-gpu=no "$nnet" "$feat" ark:- \| \
	subsample-feats --n=10 ark:- \
     ark,scp:${dir}/dvector.JOB.ark,${dir}/dvector.JOB.scp || exit 1;
fi

if [ $stage -le 1 ]; then
  echo "$0: combining dvectors across jobs"
  for j in $(seq $nj); do cat $dir/dvector.$j.scp; done >$dir/dvector.scp || exit 1;
fi

#if [ $stage -le 2 ]; then
#  # Average the utterance-level dvectors to get speaker-level dvectors.
#  echo "$0: computing mean of dvectors for each speaker"
#  $cmd $dir/log/speaker_mean.log \
#    ivector-mean ark:$data/spk2utt scp:$dir/dvector.scp ark:- ark,t:$dir/num_utts.ark \| \
#    ivector-normalize-length ark:- ark,scp:$dir/spk_dvector.ark,$dir/spk_dvector.scp || exit 1;
#fi

